{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d72971a106986eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "309a955c1342292a",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dir = os.path.expanduser(\"~/dataset/dbl/AerialPhotography\")\n",
    "output_path = os.path.expanduser(\"~/dataset/dbl/parsed_from_exif-sift_matching\")\n",
    "prior_poses = np.load(os.path.join(output_path, \"parsed.npy\"), allow_pickle=True).item()  # produced by `dji2pose.ipynb`\n",
    "prior_poses.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c7ef0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# conver arg list to command\n",
    "import shlex\n",
    "def print_command(args: list, return_str: bool = False):\n",
    "    escaped_args = []\n",
    "    for i in args:\n",
    "        escaped_args.append(shlex.quote(i))\n",
    "    connector = \" \"\n",
    "    if len(args) > 3:\n",
    "        connector = \" \\\\\\n    \"\n",
    "\n",
    "    command_str = connector.join(escaped_args)\n",
    "    if return_str:\n",
    "        return command_str\n",
    "    print(command_str)\n",
    "\n",
    "def get_command(args: list) -> str:\n",
    "    return print_command(args, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f023c4e59690e5a1",
   "metadata": {},
   "source": [
    "# 1. Distance based filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94126ed32a2ca60a",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_centers = prior_poses[\"c2w\"][:, :3, 3]\n",
    "camera_centers.shape, prior_poses[\"c2w\"][0], camera_centers[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dfe56ee3d0ec522",
   "metadata": {},
   "outputs": [],
   "source": [
    "(camera_centers[:, None, :] - camera_centers[None, :, :]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "770b06f99df6b430",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_distances = np.sqrt(np.sum(np.power(camera_centers[:, None, :] - camera_centers[None, :, :], 2), axis=-1))\n",
    "camera_distances.shape, camera_distances[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adb78b2c2efb285b",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_distances[:16, :16]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65c3be53b6610c38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make sure the calculation is correct\n",
    "for i in range(16):\n",
    "    for j in range(16):\n",
    "        assert np.isclose(np.linalg.norm(camera_centers[i] - camera_centers[j]), camera_distances[i, j]), (i, j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60d3093bfa55dda7",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_distance = 384.\n",
    "\n",
    "distance_filter_mask = camera_distances < max_distance\n",
    "distance_filter_image_count = distance_filter_mask.sum(axis=-1)\n",
    "distance_filter_image_count, \\\n",
    "    distance_filter_image_count.mean().astype(np.int32), \\\n",
    "    distance_filter_image_count.min(), \\\n",
    "    distance_filter_image_count.max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a04ec00ecd463781",
   "metadata": {},
   "source": [
    "# 2. View direction based filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20de33b4437582fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Z+ is the view direction\n",
    "camera_view_directions = prior_poses[\"c2w\"][:, :3, 2]\n",
    "camera_view_direction_cosine = np.sum(camera_view_directions[:, None, :] * camera_view_directions[None, :, :], axis=-1)\n",
    "camera_view_direction_cosine[:16, :16]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80b653da3730271c",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_view_direction_cosine.min(), camera_view_direction_cosine.max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b86d276db04b682a",
   "metadata": {},
   "source": [
    "## 2.1 Simple result validation\n",
    "Ignore this section when using other dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dda347a12d2625dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_names = prior_poses[\"image_name_list\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3b123cf4737591b",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_i = 5\n",
    "image_names[a_i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cda4ecd3ba466c30",
   "metadata": {},
   "source": [
    "Pick another image in opposite direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97e9134bdf814b38",
   "metadata": {},
   "outputs": [],
   "source": [
    "b_i = None\n",
    "for idx, i in enumerate(image_names):\n",
    "    if i == image_names[a_i].replace(\"h\", \"q\"):\n",
    "        b_i = idx\n",
    "        break\n",
    "image_names[b_i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd17723c1c30201",
   "metadata": {},
   "outputs": [],
   "source": [
    "# < 0. is expected\n",
    "camera_view_direction_cosine[a_i, b_i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e097fbbf4cca874a",
   "metadata": {},
   "source": [
    "The direction of the `x` set should be closed to -z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3edb9d92dff5c1e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_set_image_idx = 5006\n",
    "image_names[z_set_image_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72b0a4b751993755",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the z close to -Z is expected\n",
    "camera_view_directions[z_set_image_idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1521ac6a42200e50",
   "metadata": {},
   "source": [
    "The cosine value between `x` set and others should > 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "105be0aadd26d2b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, i in enumerate(image_names):\n",
    "    if os.path.basename(i).endswith(image_names[z_set_image_idx][image_names[z_set_image_idx].rfind(\"x\")+1:]):\n",
    "        print(\"{} -> {}\".format(i, camera_view_direction_cosine[z_set_image_idx, idx]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68ac3bdaef9e5fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# min < 0 may appear, but the mean should be large\n",
    "camera_view_direction_cosine[z_set_image_idx].max(), \\\n",
    "    camera_view_direction_cosine[z_set_image_idx].min(), \\\n",
    "    camera_view_direction_cosine[z_set_image_idx].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b77bf7a17e441b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# > 0 and < 1 is expected\n",
    "camera_view_direction_cosine[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"z/22z01346.JPG\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a246741057337d57",
   "metadata": {},
   "source": [
    "# 3. Z distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb927b26d091b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "w2c = np.linalg.inv(prior_poses[\"c2w\"])\n",
    "np.concatenate([prior_poses[\"c2w\"][0], w2c[0]], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4347c0c2a6ffd421",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_z_distances = np.sum(w2c[:, 2, :3][:, None, :] * camera_centers[None, :, :], axis=-1) + w2c[:, 2, 3][:, None]\n",
    "camera_z_distances[:16, :16]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc0d2ccf8d37cd00",
   "metadata": {},
   "source": [
    "## 3.1 Simple result validation\n",
    "Ignore this section when using other dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef8e428cc19380a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1346 and 1345 is in front of 1347, while 1348 behind it\n",
    "# [> 0., > 0., > 0., < 0.]\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"h/22h01346.JPG\")], \\\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"z/22z01346.JPG\")], \\\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"h/22h01345.JPG\")], \\\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"h/22h01348.JPG\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0c4736c4e41a16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# < 0.\n",
    "camera_z_distances[image_names.index(\"h/22h01346.JPG\"), image_names.index(\"h/22h01347.JPG\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4a85b50587d3c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# same position, close to 0.\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"q/22q01347.JPG\")], \\\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"x/22x01347.JPG\")], \\\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"y/22y01347.JPG\")], \\\n",
    "camera_z_distances[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"z/22z01347.JPG\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa183006dc8fc60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# in front of the camera\n",
    "camera_z_distances_mask = camera_z_distances > 0.\n",
    "camera_z_distances_mask.sum(axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "359978363a484ff5",
   "metadata": {},
   "source": [
    "# 4. Build feature matching mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ffa6da4a60a4e60",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = np.copy(distance_filter_mask)\n",
    "mask.sum(axis=-1).min()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec89fc7c3df4e1fe",
   "metadata": {},
   "source": [
    "Filter out those behind cameras, and do not have similar view direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc2eb8e1b3f78ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_cosine_value = np.cos(np.pi / 3)\n",
    "min_cosine_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa824f69b656b3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "is_behind_camera = camera_z_distances < 0.\n",
    "is_not_similar_direction = camera_view_direction_cosine < min_cosine_value\n",
    "behind_camera_filter_mask = np.logical_and(is_behind_camera, is_not_similar_direction)\n",
    "mask = np.logical_and(mask, np.logical_not(behind_camera_filter_mask))  # true represents the valid pair\n",
    "final_image_count = mask.sum(axis=-1)\n",
    "final_image_count.mean().astype(np.int32), \\\n",
    "final_image_count.min(), \\\n",
    "final_image_count.max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e478dbfe63833755",
   "metadata": {},
   "source": [
    "# 4.1 Simple result validation\n",
    "Ignore this section when using other dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ecc0561fc7d4a0",
   "metadata": {},
   "source": [
    "neighbourhoods, same view direction, should always be selected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "374f6a54a5735833",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in [45, 46]:\n",
    "    print(mask[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"h/22h013{}.JPG\".format(i))])\n",
    "    print(mask[image_names.index(\"h/22h013{}.JPG\".format(i)), image_names.index(\"h/22h01347.JPG\".format(i))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7df81a379deca73",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in [45, 46, 48, 49]:\n",
    "    print(mask[image_names.index(\"x/22x01347.JPG\"), image_names.index(\"x/22x013{}.JPG\".format(i))])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "673731b90310c49d",
   "metadata": {},
   "source": [
    "neighbourhoods, behind, opposite direction, should not be selected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "328f50d5fccb31af",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in [48, 49]:\n",
    "    print(mask[image_names.index(\"h/22h01347.JPG\"), image_names.index(\"q/22q013{}.JPG\".format(i))])\n",
    "    print(mask[image_names.index(\"q/22q013{}.JPG\".format(i)), image_names.index(\"h/22h01347.JPG\".format(i))])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "671acd1cbf8351e3",
   "metadata": {},
   "source": [
    "neighbourhoods, front, opposite direction, should be selected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "770801659d4ca461",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in [48, 49]:\n",
    "    print(mask[image_names.index(\"q/22q01347.JPG\"), image_names.index(\"h/22h013{}.JPG\".format(i))])\n",
    "    print(mask[image_names.index(\"h/22h013{}.JPG\".format(i)), image_names.index(\"q/22q01347.JPG\".format(i))])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bc24fe61af5cd39",
   "metadata": {},
   "source": [
    "# 5. Build pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e555def6207e1e58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hloc import (\n",
    "    extract_features,\n",
    "    match_features,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dc2512bb6344aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfm_pairs = os.path.join(output_path, \"pairs-netvlad.txt\")\n",
    "sfm_dir = os.path.join(output_path, \"sfm_superpoint+superglue\")\n",
    "os.makedirs(output_path, exist_ok=True)\n",
    "sfm_pairs, sfm_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9e25dab099e4e86",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieval_conf = extract_features.confs[\"netvlad\"]\n",
    "feature_conf = extract_features.confs[\"superpoint_aachen\"]\n",
    "matcher_conf = match_features.confs[\"superglue\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ca8c025cbb4d932",
   "metadata": {},
   "source": [
    "## 5.1. NetVLAD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d35141ef2fa7bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# It takes hours\n",
    "retrieval_path = extract_features.main(retrieval_conf, image_dir, output_path)\n",
    "retrieval_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7a812c59d2c65e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hloc.pairs_from_retrieval import get_descriptors\n",
    "query_desc = get_descriptors(image_names, retrieval_path)\n",
    "query_desc.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b8617aa90ea495",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = torch.einsum(\"id,jd->ij\", query_desc, query_desc)\n",
    "sim.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa736968952fcb9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hloc.pairs_from_retrieval import pairs_from_score_matrix\n",
    "num_matched = 64\n",
    "pairs = pairs_from_score_matrix(\n",
    "    torch.clone(sim),\n",
    "    np.logical_not(np.logical_and(\n",
    "        np.logical_not(np.eye(mask.shape[0], dtype=np.bool_)),  # avoid self pairing\n",
    "        mask,\n",
    "    )),\n",
    "    num_matched,\n",
    "    min_score=0.25,\n",
    ")\n",
    "# pairs = pairs_from_score_matrix(torch.clone(sim), torch.eye(sim.shape[0], dtype=torch.bool).numpy(), num_matched, min_score=0.25)\n",
    "pairs = [(image_names[i], image_names[j]) for i, j in pairs]\n",
    "pairs, len(pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b715c593d8f45ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(sfm_pairs, \"w\") as f:\n",
    "    f.write(\"\\n\".join(\" \".join([i, j]) for i, j in pairs))\n",
    "sfm_pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e23e11b342eb5cb",
   "metadata": {},
   "source": [
    "# 6. Feature matching\n",
    "Assuming all the images and cameras have been imported by `dji2pose.ipynb`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6d7c219ebbf784",
   "metadata": {},
   "outputs": [],
   "source": [
    "colmap_db_path = os.path.join(output_path, \"colmap.db\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d645a027",
   "metadata": {},
   "source": [
    "## 6.1. [Option 1] SIFT (Recommeneded)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fd47992",
   "metadata": {},
   "source": [
    "6.1.1. Extract features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2ed4a4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_command([\n",
    "    \"colmap\",\n",
    "    \"feature_extractor\",\n",
    "    \"--image_path={}\".format(image_dir),\n",
    "    \"--database_path={}\".format(colmap_db_path),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7414483d",
   "metadata": {},
   "source": [
    "6.1.2. Match features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b66aee07",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_command([\n",
    "    \"colmap\",\n",
    "    \"matches_importer\",\n",
    "    \"--database_path={}\".format(colmap_db_path),\n",
    "    \"--match_list_path={}\".format(sfm_pairs),\n",
    "    \"--match_type=pairs\",\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "625b7476",
   "metadata": {},
   "source": [
    "## 6.2. [Option 2] SuperPoint + SuperGlue"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87119323",
   "metadata": {},
   "source": [
    "6.2.1. Extract features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef9af2ef5e5c5b06",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_path = extract_features.main(feature_conf, image_dir, output_path)\n",
    "feature_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56770edc",
   "metadata": {},
   "source": [
    "6.2.2. Match features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccafb149c47351a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# It takes hours\n",
    "# LightGlue should be faster a little bit\n",
    "from pathlib import Path\n",
    "match_path = match_features.main(\n",
    "    matcher_conf, Path(sfm_pairs), feature_conf[\"output\"], output_path\n",
    ")\n",
    "match_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c279eff94fedd81c",
   "metadata": {},
   "source": [
    "6.2.3. Import to colmap database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38337323232d446f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hloc.reconstruction import (\n",
    "    import_features,\n",
    "    import_matches,\n",
    "    estimation_and_geometric_verification,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7991dcafa0226df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# map `image_name` to colmap `image_id`\n",
    "import sqlite3\n",
    "db = sqlite3.connect(colmap_db_path)\n",
    "try:\n",
    "    image_ids = {}\n",
    "    for name, image_id in db.execute(\"SELECT name, image_id FROM images;\"):\n",
    "        image_ids[name] = image_id\n",
    "finally:\n",
    "    db.close()\n",
    "image_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7a55d7a9c965226",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clear matchs, keypoints and two view geometries\n",
    "# db = sqlite3.connect(colmap_db_path)\n",
    "# try:\n",
    "#     db.execute(\"DELETE FROM `matches`\")\n",
    "#     db.execute(\"DELETE FROM `keypoints`\")\n",
    "#     db.execute(\"DELETE FROM `two_view_geometries`\")\n",
    "#     db.commit()\n",
    "# finally:\n",
    "#     db.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5277ee96c34b82c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import_features(image_ids, colmap_db_path, feature_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "392fa1c284db8772",
   "metadata": {},
   "outputs": [],
   "source": [
    "import_matches(\n",
    "    image_ids,\n",
    "    colmap_db_path,\n",
    "    sfm_pairs,\n",
    "    match_path,\n",
    "    None,\n",
    "    False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c5965ee2269411",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimation_and_geometric_verification(colmap_db_path, sfm_pairs, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9096c9f7",
   "metadata": {},
   "source": [
    "# 7. Mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78bbe6cd",
   "metadata": {},
   "source": [
    "## 7.1. [Option 1] colmap hierarchical_mapper\n",
    "First choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924316e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sparse_hierarchical_mapper_dir = os.path.join(output_path, \"sparse-hierarchical_mapper\")\n",
    "os.makedirs(sparse_hierarchical_mapper_dir, exist_ok=True)\n",
    "print(\" \\\\\\n    \".join([\n",
    "    \"colmap\",\n",
    "    \"hierarchical_mapper\",\n",
    "    \"--database_path={}\".format(colmap_db_path),\n",
    "    \"--image_path={}\".format(image_dir),\n",
    "    \"--output_path={}\".format(sparse_hierarchical_mapper_dir),\n",
    "    \"--image_overlap=100\",\n",
    "    \"--leaf_max_num_images=900\",\n",
    "    \"--Mapper.ba_use_gpu=1\",\n",
    "]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32093096",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run a `point_triangulator` is recommended\n",
    "largest_sparse_model_dir = os.path.join(sparse_hierarchical_mapper_dir, \"0\")  # Remember to change this\n",
    "largest_sparse_model_retri_output_dir = \"{}-retriangulated\".format(largest_sparse_model_dir)\n",
    "os.makedirs(largest_sparse_model_retri_output_dir, exist_ok=True)\n",
    "print(\" \\\\\\n    \".join([\n",
    "    \"colmap\",\n",
    "    \"point_triangulator\",\n",
    "    \"--database_path={}\".format(colmap_db_path),\n",
    "    \"--image_path={}\".format(image_dir),\n",
    "    \"--input_path={}\".format(largest_sparse_model_dir),\n",
    "    \"--output_path={}\".format(largest_sparse_model_retri_output_dir),\n",
    "    \"--Mapper.ba_use_gpu=1\",\n",
    "]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b2906f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run a `bundle_adjuster` is recommended, but outliers should be removed first.\n",
    "# Take a look `gps_based_sfm_outlier_detection.ipynb`\n",
    "largest_sparse_model_retri_ba_output_dir = \"{}-bundle_adjusted\".format(largest_sparse_model_retri_output_dir)\n",
    "os.makedirs(largest_sparse_model_retri_ba_output_dir, exist_ok=True)\n",
    "print_command([\n",
    "    \"colmap\",\n",
    "    \"bundle_adjuster\",\n",
    "    \"--input_path={}\".format(largest_sparse_model_retri_output_dir),\n",
    "    \"--output_path={}\".format(largest_sparse_model_retri_ba_output_dir),\n",
    "    \"--BundleAdjustment.use_gpu=1\",\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f24df32",
   "metadata": {},
   "source": [
    "## 7.2. [Option 2] Custom partitioning\n",
    "Experimental"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "388a58e6c8bad6a6",
   "metadata": {},
   "source": [
    "### 7.2.1. Partitioning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa796e11f6f4d44",
   "metadata": {},
   "outputs": [],
   "source": [
    "from internal.utils.partitioning_utils import SceneConfig, PartitionableScene\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7398bd26504535c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_centers = prior_poses[\"c2w\"][:, :3, 3]\n",
    "reoriented_camera_centers = torch.from_numpy(camera_centers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95c4b48bd5bfc426",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect('equal', adjustable='box')\n",
    "scene_size = torch.max(reoriented_camera_centers, dim=0).values - torch.min(reoriented_camera_centers, dim=0).values\n",
    "ax.set_xlim([torch.min(reoriented_camera_centers[:, 0]) - 0.1 * scene_size[0], torch.max(reoriented_camera_centers[:, 0]) + 0.1 * scene_size[0]])\n",
    "ax.set_ylim([torch.min(reoriented_camera_centers[:, 1]) - 0.1 * scene_size[1], torch.max(reoriented_camera_centers[:, 1]) + 0.1 * scene_size[1]])\n",
    "ax.scatter(reoriented_camera_centers[:, 0], reoriented_camera_centers[:, 1], s=0.2, c=\"red\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24bc5d0dece27f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_config = SceneConfig(\n",
    "    origin=torch.tensor([0., 0.]),\n",
    "    partition_size=500.,\n",
    ")\n",
    "scene = PartitionableScene(scene_config, reoriented_camera_centers[..., :2])\n",
    "scene.get_bounding_box_by_camera_centers()\n",
    "scene.get_scene_bounding_box()\n",
    "scene.plot(scene.plot_scene_bounding_box)\n",
    "scene.build_partition_coordinates()\n",
    "scene.plot(scene.plot_partitions)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207db12409235e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_config.location_based_enlarge = 0.\n",
    "assignment_without_enlarging = scene.camera_center_based_partition_assignment()\n",
    "assignment_without_enlarging.sum(-1).reshape(scene.scene_bounding_box.n_partitions.tolist()[::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16d16893ae50e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# merge partition with image number less than `min_images` into its most visible neighbour\n",
    "min_images = 255\n",
    "\n",
    "assignment_without_enlarging_2d = assignment_without_enlarging.reshape((scene.scene_bounding_box.n_partitions.tolist()[::-1] + [-1]))\n",
    "\n",
    "assignment_merged = torch.clone(assignment_without_enlarging_2d)\n",
    "\n",
    "def get_number_of_images(x, y) -> int:\n",
    "    if x < 0:\n",
    "        return 0\n",
    "    if y < 0:\n",
    "        return 0\n",
    "    if x >= scene.scene_bounding_box.n_partitions[0]:\n",
    "        return 0\n",
    "    if y >= scene.scene_bounding_box.n_partitions[1]:\n",
    "        return 0\n",
    "    return assignment_merged[y, x].sum().item()\n",
    "\n",
    "for y in range(scene.scene_bounding_box.n_partitions[1]):\n",
    "    for x in range(scene.scene_bounding_box.n_partitions[0]):\n",
    "        n_images = get_number_of_images(x, y)\n",
    "        if n_images == 0:\n",
    "            continue\n",
    "        if n_images > min_images:\n",
    "            continue\n",
    "            \n",
    "        partition_camera_mean_center = camera_centers[assignment_merged[y, x]].mean(axis=0)\n",
    "            \n",
    "        max_neighbour_n_visible = -1\n",
    "        merge_to_neighbour = None\n",
    "        \n",
    "        for neighbour in [(x-1, y), (x+1, y), (x, y-1), (x, y+1)]:\n",
    "            neighbour_n_images = get_number_of_images(*neighbour)\n",
    "            if neighbour_n_images == 0:\n",
    "                continue\n",
    "                \n",
    "            partition_image_visible_masks = torch.from_numpy(mask[assignment_merged[y, x].numpy()])  # [N_local_images, N_images]\n",
    "            partition_neighbour_image_visible_mask = partition_image_visible_masks[:, assignment_without_enlarging_2d[neighbour[1], neighbour[0]]]  # [N_local_images, N_neighbour_partition_images]\n",
    "            n_visible_neighbour = partition_neighbour_image_visible_mask.sum()\n",
    "            \n",
    "            if n_visible_neighbour > max_neighbour_n_visible:\n",
    "                max_neighbour_n_visible = n_visible_neighbour\n",
    "                merge_to_neighbour = neighbour\n",
    "                \n",
    "            # neighbour_distance = np.linalg.norm(camera_centers[assignment_without_enlarging_2d[y, x]] - partition_camera_mean_center[None, :], axis=-1).min()\n",
    "            # print(neighbour, neighbour_distance)\n",
    "            #     \n",
    "            # if neighbour_distance < min_neighbour_distance:\n",
    "            #     min_neighbour_n_images = neighbour_n_images\n",
    "            #     min_neighbour_distance = neighbour_distance\n",
    "            #     min_neighbour = neighbour\n",
    "                \n",
    "        assert merge_to_neighbour is not None, (x, y)\n",
    "        # assert min_neighbour_n_images > 0, (x, y)\n",
    "        \n",
    "        assignment_merged[merge_to_neighbour[1], merge_to_neighbour[0]] = torch.logical_or(\n",
    "            assignment_merged[y, x],\n",
    "            assignment_merged[merge_to_neighbour[1], merge_to_neighbour[0]],\n",
    "        )\n",
    "        assignment_merged[y, x].fill_(False)\n",
    "        \n",
    "        print(\"({}, {}) -> ({}, {})\".format(x + scene.scene_bounding_box.origin_partition_offset[0], y + scene.scene_bounding_box.origin_partition_offset[1], merge_to_neighbour[0] + scene.scene_bounding_box.origin_partition_offset[0], merge_to_neighbour[1] + scene.scene_bounding_box.origin_partition_offset[1]))\n",
    "        # break\n",
    "    # break\n",
    "assignment_merged.sum(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "811ff95f71750921",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assign with enlarged partition\n",
    "scene_config.location_based_enlarge = 0.2  # the merging requires overlap\n",
    "scene.camera_center_based_partition_assignment().sum(-1).reshape(scene.scene_bounding_box.n_partitions.tolist()[::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bac8056aa933764b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mask out those partitions merged into others\n",
    "masked_assignment = torch.logical_and(scene.is_camera_in_partition, assignment_merged.sum(-1).reshape((-1, 1)) > 0)\n",
    "masked_assignment.sum(-1).reshape(scene.scene_bounding_box.n_partitions.tolist()[::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f7b366d26e56b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# merge the enlarged and merged assignments\n",
    "final_assignment = torch.logical_or(masked_assignment, assignment_merged.reshape(masked_assignment.shape))\n",
    "final_assignment.sum(-1).reshape(scene.scene_bounding_box.n_partitions.tolist()[::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13fd2d718bd0d067",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import cm\n",
    "import matplotlib.patches as mpatches\n",
    "import random\n",
    "\n",
    "plt.close()\n",
    "fig, ax = plt.subplots()\n",
    "scene.set_plot_ax_limit(ax)\n",
    "ax.set_aspect('equal', adjustable='box')\n",
    "colors = list(iter(cm.rainbow(np.linspace(0, 1, len(scene.partition_coordinates)))))\n",
    "random.shuffle(colors)\n",
    "color_iter = iter(colors)\n",
    "\n",
    "annotate_position_x = 0.125\n",
    "annotate_position_y = 0.25\n",
    "annotate_font_size = 5\n",
    "\n",
    "idx = 0\n",
    "local_idx = 0\n",
    "for partition_id, partition_xy in scene.partition_coordinates:\n",
    "    try:\n",
    "        assigned_camera_mask = scene.is_camera_in_partition[idx]\n",
    "        color=next(color_iter)\n",
    "        \n",
    "        if assigned_camera_mask.sum() == 0:\n",
    "            continue\n",
    "            \n",
    "        assigned_camera_centers = camera_centers[assigned_camera_mask.numpy()]\n",
    "                        \n",
    "        ax.scatter(\n",
    "            assigned_camera_centers[:, 0], \n",
    "            assigned_camera_centers[:, 1],\n",
    "            s=0.2,\n",
    "            c=color[None, :] * np.asarray([[1., 1., 1., 0.1]]),\n",
    "        )\n",
    "\n",
    "        ax.add_artist(mpatches.Rectangle(\n",
    "            (partition_xy[0], partition_xy[1]),\n",
    "            scene.scene_config.partition_size,\n",
    "            scene.scene_config.partition_size,\n",
    "            fill=False,\n",
    "            color=color,\n",
    "        ))\n",
    "        ax.annotate(\n",
    "            \"#{}\\n({}, {})\".format(local_idx, partition_id[0], partition_id[1]),\n",
    "            xy=(\n",
    "                partition_xy[0] + annotate_position_x * scene.scene_config.partition_size,\n",
    "                partition_xy[1] + annotate_position_y * scene.scene_config.partition_size,\n",
    "            ),\n",
    "            fontsize=annotate_font_size,\n",
    "        )\n",
    "        local_idx += 1\n",
    "    finally:\n",
    "        idx += 1\n",
    "\n",
    "plt.savefig(os.path.join(output_path, \"sfm_partitions.png\"), dpi=600)\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6cfcec1fb9184f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make sure all images are covered\n",
    "assert (final_assignment.sum(dim=0) > 0).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44cbe6b6b88e8521",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(os.path.join(output_path, \"sfm_partitions.npy\"), final_assignment.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f195be8d5761425",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_assignment.numpy().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb56d9ad9f6eaa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_name_array = np.asarray(image_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f02bd6d41639e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_list_output_path = os.path.join(output_path, \"sfm_partition_image_lists\")\n",
    "os.makedirs(image_list_output_path, exist_ok=True)\n",
    "\n",
    "for i in os.scandir(image_list_output_path):\n",
    "    if not i.is_dir(follow_symlinks=False):\n",
    "        os.unlink(i.path)\n",
    "\n",
    "local_idx = 0\n",
    "for partition in final_assignment:\n",
    "    if partition.sum() == 0:\n",
    "        continue\n",
    "        \n",
    "    partition_image_name_array = image_name_array[partition]\n",
    "    with open(os.path.join(image_list_output_path, \"{:02d}.txt\".format(local_idx)), \"w\") as f:\n",
    "        for i in partition_image_name_array:\n",
    "            f.write(i)\n",
    "            f.write(\"\\n\")\n",
    "    \n",
    "    local_idx += 1\n",
    "n_sfm_partitions = local_idx\n",
    "image_list_output_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5745a4906554ebdf",
   "metadata": {},
   "source": [
    "### 7.2.2. Mapping\n",
    "Build colmap with cuDSS to enable mapper with GPU\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc5de17b8f902a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfm_partition_dir = os.path.join(output_path, \"sfm_partitions\")\n",
    "os.makedirs(sfm_partition_dir, exist_ok=True)\n",
    "\n",
    "remove_abs_path = len(os.path.dirname(image_dir)) + 1\n",
    "\n",
    "print_command([\"chmod\", \"400\", colmap_db_path[remove_abs_path:]])  # key to run multiple mapper in parallel\n",
    "print()\n",
    "for partition_image_list_file in os.scandir(os.path.join(output_path, \"sfm_partition_image_lists\")):\n",
    "    partition_idx = int(partition_image_list_file.name.rsplit(\".\")[0])\n",
    "    sfm_partition_sparse_dir = os.path.join(sfm_partition_dir, \"{:02d}\".format(partition_idx))\n",
    "    os.makedirs(sfm_partition_sparse_dir, exist_ok=True)\n",
    "    print_command([\"mkdir\", \"-p\", sfm_partition_sparse_dir[remove_abs_path:]])\n",
    "    print(get_command([\n",
    "        \"srun\",\n",
    "        \"--gpus=1\",\n",
    "        \"--nodes=1\",\n",
    "        \"--ntasks=1\",\n",
    "        \"--exclusive\",\n",
    "        \"colmap\",\n",
    "        \"mapper\",\n",
    "        \"--image_list_path={}\".format(partition_image_list_file.path[remove_abs_path:]),\n",
    "        \"--Mapper.ba_use_gpu=1\",\n",
    "        # \"--Mapper.min_num_matches=32\",\n",
    "        \"--database_path={}\".format(colmap_db_path[remove_abs_path:]),\n",
    "        \"--image_path={}\".format(image_dir[remove_abs_path:]),\n",
    "        \"--output_path={}\".format(sfm_partition_sparse_dir[remove_abs_path:]),        \n",
    "    ]) + \" &\\n\")\n",
    "print(\"wait\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374129014310a362",
   "metadata": {},
   "source": [
    "### 7.2.3. Remove cameras with large error values compared to GPS\n",
    "This section can be run independently"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "328ed1511a01a4e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "from internal.utils import sfm_outlier_detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8ffabdf5622faa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfm_partition_load_dir = os.path.expanduser(\"~/dataset/dbl/parsed_from_exif-sift_matching/sfm_partitions\")\n",
    "# outlier_removed_partition_output_dir = os.path.expanduser(\"~/dataset/dbl/parsed_from_exif-sift_matching/sfm_partitions-outliers_removed\")\n",
    "outlier_removed_partition_output_dir = \"{}-outliers_removed\".format(sfm_partition_load_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "283efbea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_largest_sparse_model(path: str):\n",
    "    # find model with the most images\n",
    "    sfm_partition_sparse_model_dir = None\n",
    "    sfm_partition_cameras_bin_max_size = -1\n",
    "    for sparse_model_idx_dir in os.scandir(path):\n",
    "        if not sparse_model_idx_dir.is_dir():\n",
    "            continue\n",
    "        images_bin_path = os.path.join(sparse_model_idx_dir.path, \"images.bin\")\n",
    "        if not os.path.exists(images_bin_path):\n",
    "            continue\n",
    "        images_bin_size = os.stat(images_bin_path).st_size\n",
    "        if images_bin_size > sfm_partition_cameras_bin_max_size:\n",
    "            sfm_partition_cameras_bin_max_size = images_bin_size\n",
    "            sfm_partition_sparse_model_dir = sparse_model_idx_dir.path\n",
    "    assert sfm_partition_sparse_model_dir is not None\n",
    "    return sfm_partition_sparse_model_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8fa2519",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sfm_partition_dir in tqdm(sorted(list(os.scandir(sfm_partition_load_dir)), key=lambda i: i.name)):\n",
    "    if not sfm_partition_dir.is_dir():\n",
    "        continue\n",
    "    sfm_partition_idx_str = sfm_partition_dir.name\n",
    "    sfm_partition_idx = int(sfm_partition_idx_str)\n",
    "\n",
    "    largest_sparse_model_dir = find_largest_sparse_model(sfm_partition_dir.path)\n",
    "    \n",
    "    colmap_images = sfm_outlier_detection.load(prior_poses, largest_sparse_model_dir)\n",
    "    sfm_outlier_detection.filter(\n",
    "        colmap_images,\n",
    "        output_path=os.path.join(outlier_removed_partition_output_dir, sfm_partition_idx_str),\n",
    "        min_acceptable_error_limit=6.,\n",
    "        max_acceptable_error_limit=64.,\n",
    "    )()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84b10de6",
   "metadata": {},
   "source": [
    "### 7.2.4. Merge partitions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "599fbdac",
   "metadata": {},
   "source": [
    "Define the merge order\n",
    "\n",
    "TODO: calculate the order based on overlapping\n",
    "\n",
    "TODO: BA for every N partitions merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5e359a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Partition IDs; DFS\n",
    "merge_order = [\n",
    "    [1, 2, 0],\n",
    "    [3, 4, 5],\n",
    "    [6, 7, 8, 9],\n",
    "    [10, 11, 12, 13],\n",
    "    [14, 15],\n",
    "]\n",
    "merged_partition_output_path = \"{}-merged\".format(sfm_partition_load_dir)\n",
    "merged_partition_output_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95e67d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union, List\n",
    "import subprocess\n",
    "\n",
    "# validate the merge order\n",
    "partition_id_appear_counters = {}\n",
    "def validate_merge_order(m):\n",
    "    if isinstance(m, List) and len(m) == 1:\n",
    "        m = m[0]\n",
    "    if isinstance(m, int):\n",
    "        assert m not in partition_id_appear_counters, \"partition #{} appears twice\".format(m)\n",
    "        partition_id_appear_counters[m] = 1\n",
    "        return 1\n",
    "\n",
    "    total = 0\n",
    "    for i in m:\n",
    "        total += validate_merge_order(i)\n",
    "    return total\n",
    "\n",
    "n_partitions_to_merge = validate_merge_order(merge_order)\n",
    "merged_counter = 1\n",
    "\n",
    "# DFS merging\n",
    "def merge_partitions(partition_ids: Union[int, List]) -> str:\n",
    "    if isinstance(partition_ids, List) and len(partition_ids) == 1:\n",
    "        partition_ids = partition_ids[0]\n",
    "    if isinstance(partition_ids, int):\n",
    "        return os.path.join(outlier_removed_partition_output_dir, \"{:02d}\".format(partition_ids))\n",
    "    \n",
    "    sub_partitions = []\n",
    "    for i in partition_ids:\n",
    "        sub_partitions.append(merge_partitions(i))\n",
    "\n",
    "    merged = sub_partitions[0]\n",
    "    for i in sub_partitions[1:]:\n",
    "        output_name = \"{}-{}\".format(os.path.basename(merged), os.path.basename(i))\n",
    "        output_path = os.path.join(merged_partition_output_path, output_name)\n",
    "        os.makedirs(output_path, exist_ok=True)\n",
    "        if os.path.exists(os.path.join(output_path, \"images.bin\")):\n",
    "            print(\"skip {}\".format(output_name))\n",
    "        else:\n",
    "            args = [\n",
    "                \"colmap\",\n",
    "                \"model_merger\",\n",
    "                \"--input_path1={}\".format(merged),\n",
    "                \"--input_path2={}\".format(i),\n",
    "                \"--output_path={}\".format(output_path),\n",
    "            ]\n",
    "            print_command(args)\n",
    "            assert subprocess.call(args) == 0\n",
    "        global merged_counter\n",
    "        merged_counter += 1\n",
    "        print(\"{} of {} merged\".format(merged_counter, n_partitions_to_merge))\n",
    "        merged = output_path\n",
    "    \n",
    "    return merged\n",
    "\n",
    "    \n",
    "final_merged_sparse_model = merge_partitions(merge_order)\n",
    "final_merged_sparse_model_alias = os.path.join(os.path.dirname(final_merged_sparse_model), \"final\")\n",
    "if os.path.exists(final_merged_sparse_model_alias):\n",
    "    os.unlink(final_merged_sparse_model_alias)\n",
    "os.symlink(final_merged_sparse_model, final_merged_sparse_model_alias)\n",
    "final_merged_sparse_model_alias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "738c6f47",
   "metadata": {},
   "source": [
    "### 7.2.5. point_triangulator and bundle_adjuster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0306b3bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "tri_output_path = \"{}-retriangulated\".format(final_merged_sparse_model_alias)\n",
    "os.makedirs(tri_output_path, exist_ok=True)\n",
    "print_command([\n",
    "    \"colmap\",\n",
    "    \"point_triangulator\",\n",
    "    \"--database_path={}\".format(colmap_db_path),\n",
    "    \"--image_path={}\".format(image_dir),\n",
    "    \"--input_path={}\".format(final_merged_sparse_model_alias),\n",
    "    \"--output_path={}\".format(tri_output_path),\n",
    "    \"--Mapper.ba_use_gpu=1\",\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6159826d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turn to `gps_based_sfm_outlier_detection.ipynb`, or BA directly.\n",
    "\n",
    "ba_output_path = \"{}-bundle_adjusted\".format(tri_output_path)\n",
    "os.makedirs(ba_output_path, exist_ok=True)\n",
    "print_command([\n",
    "    \"colmap\",\n",
    "    \"bundle_adjuster\",\n",
    "    \"--input_path={}\".format(tri_output_path),\n",
    "    \"--output_path={}\".format(ba_output_path),\n",
    "    \"--BundleAdjustment.use_gpu=1\",\n",
    "])"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
